import random

import cv2
import keras
import numpy as np
from keras import Input
from matplotlib import pyplot as plt
from config import Config, ModelLayers
# from datamanager import data_manager
from configuration.model_congig import ModelConfig
from garbage_dataset import GarbageDataSet
from model import Model
from util.model_listutil import ModelListUtil

g1 = GarbageDataSet()
x_train, y_train, x_test, y_test = g1.load_data("data/data" + str(Config.IMAGE_HEIGHT) + ".npy")


def n4ton(data):
    count = 0
    for col in data:
        if col == 1:
            return count
        count += 1


labels = ["Harmful", "Kitchen", "Other", "Recyclable"]


def randomimg():
    idxs = []
    for i in range(4):
        for j in range(5):
            idx = int(random.random() * len(x_test) // 4)+len(x_test)//4*i
            idxs.append(idx)
    print(idxs)
    plt.figure(1, figsize=(10, 10))
    for i in range(4):

        for j in range(5):
            idx = idxs[i*5+j]
            # [i * 16000: (i + 1) * 16000]
            res = x_test[idx]
            # res = x_test[i * len(x_test) // 4 + idx]
            b, g, r = cv2.split(res)
            img = cv2.merge((r, g, b))
            plt.subplot(4, 5, i * 5 + j + 1)
            # plt.subplot(4, 5, i * 5 + j + 1, xlabel=labels[n4ton(y_test[idx])])
            plt.xticks([]), plt.yticks([])
            plt.imshow(img, 'gray')

    plt.savefig("20.png", bbox_inches='tight')
    plt.show()

    random.shuffle(idxs)
    plt.figure(figsize=(10, 10))
    print(idxs)
    for i in range(4):
        for j in range(5):
            idx = idxs[i*5+j]
            res = x_test[idx]
            b, g, r = cv2.split(res)
            img = cv2.merge((r, g, b))
            plt.subplot(4, 5, i * 5 + j + 1)
            plt.xticks([]), plt.yticks([])
            plt.imshow(img, 'gray')

    plt.savefig("20notag.png", bbox_inches='tight')
    plt.show()


def maxid(l: list):
    l.index(max(l))


def testimgs(train_id: int):
    conf = ModelConfig()
    conf.IMAGE_HEIGHT = Config.IMAGE_HEIGHT
    conf.IMAGE_WIDTH = Config.IMAGE_WIDTH

    inputs = Input(shape=(conf.IMAGE_HEIGHT, conf.IMAGE_HEIGHT, conf.COLOR_CHANNELS))

    m1 = Model(conf, inputs)  # 0.9962
    model_list = ModelListUtil.get_model_list_from_db(620)
    m1.load(model_list, f"storage/{620}.h5")

    m2 = Model(conf, inputs)  # 0.9766
    model_list = ModelListUtil.get_model_list_from_db(502)
    m2.load(model_list, f"storage/{502}.h5")

    m3 = Model(conf, inputs)  # 0.9922
    model_list = ModelListUtil.get_model_list_from_db(504)
    m3.load(model_list, f"storage/{504}.h5")

    m4 = Model(conf, inputs)  # 0.9796
    model_list = ModelListUtil.get_model_list_from_db(505)
    m4.load(model_list, f"storage/{505}.h5")
    m11 = keras.Model(inputs, m1.model.layers[-1].output)
    m21 = keras.Model(inputs, m2.model.layers[-1].output)
    m31 = keras.Model(inputs, m3.model.layers[-1].output)
    m41 = keras.Model(inputs, m4.model.layers[-1].output)
    m11.trainable = False
    m21.trainable = False
    m31.trainable = False
    m41.trainable = False
    m = Model(conf, inputs)
    Avglist = [
        [ModelLayers.Average],
        # [ModelLayers.Dense, 64],
        # [ModelLayers.Dense, 32],
    ]

    m.create_by_list_real_cls_functional(Avglist, [m11, m21, m31, m41], False)
    m.compile()
    # m.model.load_weights(f"storage/{train_id}.h5")

    plt.figure(1, figsize=(6, 6))
    for i in range(4):

        for j in range(5):
            idx = int(random.random() * len(x_test) // 4)
            # [i * 16000: (i + 1) * 16000]
            res = x_test[i * len(x_test) // 4 + idx]
            b, g, r = cv2.split(res)
            img = cv2.merge((r, g, b))
            print(res.shape)
            plt.subplot(4, 5, i * 5 + j + 1, xlabel=labels[n4ton(y_test[i * len(x_test) // 4 + idx])] +
                                                    str(m3.model.evaluate(
                                                        np.array([res]), np.array([[1, 0, 0, 0]]), verbose=0)))
            plt.xticks([]), plt.yticks([])
            plt.imshow(img, 'gray')
    plt.show()


if __name__ == '__main__':
    randomimg()
    # testimgs(600)
